#include "matmul_op.h"
#include "transpose_op.h"
#include "../../factory/factory.h"
#include "../../factory/transpose_factory.h"
#include "../../factory/matmul_factory.h"

MatMul::MatMul(Node* lhs, Node* rhs) : OpNode({ lhs, rhs}){
    setOutShape(new Shape(lhs -> getOutShape() -> rows, rhs -> getOutShape() -> cols));
}

vector<Node*> MatMul::getGradNodes(Node* grad){
    auto nodes = *getNodeIn();
    Factory* factory;
    factory = new TransposeFactory(nodes[0]);
    Node* l_transpose = factory -> createNode();
    delete factory;
    factory = new TransposeFactory(nodes[1]);
    Node* r_transpose = factory -> createNode();
    delete factory;
    factory = new MatmulFactory(grad, r_transpose);
    Node* l_grad = factory -> createNode();
    delete factory;
    factory = new MatmulFactory(l_transpose, grad);
    Node* r_grad = factory -> createNode();
    delete factory;
    return {l_grad, r_grad};
}

void MatMul::compute(Context* context){
    auto nodes = *getNodeIn();
    Matrix l_mat = context -> getFlowedData(nodes[0]);
    Matrix r_mat = context -> getFlowedData(nodes[1]);
    context -> setFlowedData(this, matmul(l_mat, r_mat));
}
